from z3 import *
import random
import os
from openai import OpenAI
population_size = 50
num_generations = 100
mutation_rate = 0.1
s1 = """
public class Main {
    public static int test(int x, int y) {
        if (10 > x && x > 0 && x + y > 5) {
            return 1;
        } else {
            return -1;
        }
    }

    public static void main(String[] args) {
        // 示例调用
        int x = 3;
        int y = 3;
        int result = test(x, y);
        System.out.println("Result: " + result);
    }
}
"""

# s2 = """
# 将这段程序转为python语言（只返回代码即可）
# """
s2 = """
提取这段代码的约束条件,'&&'用','代替,用python中的布尔型数组存储每个约束条件,保留变量名（只返回代码即可）
"""

question = s1 + s2

client = OpenAI(
    api_key="sk-067ca2043e80470ba2313779960367f8",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)
completion = client.chat.completions.create(
    model="qwen-plus",
    # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    messages=[
        {'role': 'system', 'content': 'You are a helpful assistant.'},
        {'role': 'user', 'content': question}],
)

#解析返回的JSON数据并提取content字段
response_content = completion.choices[0].message.content
response_content = response_content.split('python')[1]
response_content = response_content.split('```')[0]
response_content = response_content.strip()
print(response_content)


# 定义变量
x = Int('x')
y = Int('y')

# 定义约束条件
exec(response_content)

print(constraints)

# 使用Z3求解器验证约束条件
s = Solver()
s.add(constraints)
is_sat = s.check()

# 如果存在解，则获取模型
solution = None
if is_sat == sat:
    m = s.model()
    solution = {
        m[x].decl().name(): m[x].as_long(),
        m[y].decl().name(): m[y].as_long()
    }

# 遗传算法参数


def generate_initial_population():
    return [{'x': random.randint(1, 15), 'y': random.randint(1, 15)} for _ in range(population_size)]

def fitness(test_case):
    # 计算适应度函数，这里简单地返回测试用例是否满足约束条件
    return int(test_case['x'] + test_case['y'] >= 10 and test_case['x'] <= test_case['y'])

def crossover(parent1, parent2):
    child = {}
    for key in parent1.keys():
        child[key] = parent1[key] if random.random() > 0.5 else parent2[key]
    return child

def mutate(test_case):
    if random.random() < mutation_rate:
        key = random.choice(list(test_case.keys()))
        test_case[key] += random.choice([-1, 1])
    return test_case

def genetic_algorithm():
    population = generate_initial_population()

    for generation in range(num_generations):
        # 选择
        population.sort(key=lambda tc: fitness(tc), reverse=True)

        # 交叉
        new_population = []
        for i in range(population_size // 2):
            parent1, parent2 = random.sample(population[:population_size // 2], 2)
            child = crossover(parent1, parent2)
            new_population.append(child)
            new_population.append(mutate(child))

        # 替换旧种群
        population = new_population

        # 输出当前世代的最佳个体
        best_individual = max(population, key=fitness)
        print(f"Generation {generation}: Best individual {best_individual} with fitness {fitness(best_individual)}")

    return best_individual

# 运行遗传算法
best_test_case = genetic_algorithm()
print(f"Best test case found: {best_test_case}")